import numpy as np
from matplotlib import pyplot as plt
from sklearn import preprocessing
from sklearn.linear_model import LinearRegression, LassoCV, RidgeCV
from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import VotingRegressor

import pandas as pd
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.svm import SVR
from sklearn.tree import DecisionTreeRegressor


def train(dataPath, modelSavePath):
    dataset = pd.read_csv(dataPath)
    X = dataset.iloc[:, :-1]
    # X = dataset.iloc[:, [0,1,2,3,5,9,10,12]]
    y = dataset.iloc[:, dataset.shape[1] - 1]

    # normalize
    scl = preprocessing.StandardScaler()
    X_scaled = scl.fit_transform(X)

    if dataPath == '../../data/ingotRate.csv':
        X_train, X_test, y_train, y_test = train_test_split(X_scaled, y.values, test_size=0.2, random_state=3)
    elif dataPath == '../../data/yieldRate.csv':
        X_train, X_test, y_train, y_test = train_test_split(X_scaled, y.values, test_size=0.2, random_state=0)

    r1 = DecisionTreeRegressor(max_depth=1, random_state=0)
    r2 = KNeighborsRegressor(n_neighbors=10)
    r3 = LassoCV()
    r4 = MLPRegressor(hidden_layer_sizes=(14, 8, 4, 2), activation='tanh', solver='lbfgs', alpha=1e-2,
                      max_iter=2000, random_state=0)
    r5 = LinearRegression()
    r6 = RidgeCV()
    r7 = SVR(kernel='linear', C=0.1, gamma='scale', degree=2, epsilon=0.01)
    regressor = VotingRegressor([('tree', r1), ('knn', r2), ('lasso', r3), ('mlp', r4), ('linear', r5), ('ridge', r6), ('svr', r7)])

    regressor.fit(X_train, y_train)

    y_train_pred = regressor.predict(X_train)

    # plt.plot(y_train, label='y_train')
    plt.plot(y_train, label='y_train')
    plt.plot(y_train_pred, label='y_train_pred')
    plt.title('EL_stacking', fontsize='large', fontweight='bold')
    plt.legend()
    plt.show()

    # print('\ntrain: ')
    # print("MSE: ", np.mean(-train_losses['MSE']))
    # print("MAE: ", np.mean(-train_losses['MAE']))
    # print("R2_Score: ", np.mean(train_losses['R2']))
    # print('\nvalid: ')
    # print("MSE: ", np.mean(-valid_losses['MSE']))
    # print("MAE: ", np.mean(-valid_losses['MAE']))
    # print("R2_Score: ", np.mean(valid_losses['R2']))

    print('\ntrain: ')
    MSE = mean_squared_error(y_train, y_train_pred)
    MAE = mean_absolute_error(y_train, y_train_pred)
    R2_Score = r2_score(y_train, y_train_pred)
    print("MSE: ", MSE)
    print("MAE: ", MAE)
    print("R2_Score: ", R2_Score)

    y_pred = regressor.predict(X_test)

    # plt.plot(y_test, label='y_test')
    plt.plot(y_test, label='y_test')
    plt.plot(y_pred, label='y_pred')
    plt.title('EL_stacking', fontsize='large', fontweight='bold')
    plt.legend()
    plt.show()

    print('\ntest: ')
    MSE = mean_squared_error(y_test, y_pred)
    MAE = mean_absolute_error(y_test, y_pred)
    R2_Score = r2_score(y_test, y_pred)
    print("MSE: ", MSE)
    print("MAE: ", MAE)
    print("R2_Score: ", R2_Score)


if __name__ == '__main__':
    train('../../data/ingotRate.csv', './ingotRatePrediction.pkl')
    # train('../../data/yieldRate.csv', './yieldRatePrediction.pkl')
